"""
Game Screens Module
Contains all screen/state classes for the game
"""

import pygame
import os
from enemy_config import get_text, ENEMY_DAMAGE


class BaseScreen:
    """
    Base class for all game screens with common functionality
    """
    
    def __init__(self, game):
        """
        Initialize base screen
        
        Args:
            game: Reference to main game object
        """
        self.game = game
        
    def handle_event(self, event):
        """
        Handle pygame events for this screen
        
        Args:
            event: Pygame event to handle
            
        Returns:
            str or None: Next screen state or None to stay on current screen
        """
        pass
        
    def update(self, dt):
        """
        Update screen logic
        
        Args:
            dt (float): Delta time since last frame
        """
        pass
        
    def draw(self, surface):
        """
        Draw screen content
        
        Args:
            surface (pygame.Surface): Surface to draw on
        """
        pass


class MenuScreen(BaseScreen):
    """
    Main menu screen with start, quit, and language options
    """
    
    def __init__(self, game):
        """Initialize menu screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle menu input events"""
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_s:
                return "instructions"  # Start game
            elif event.key == pygame.K_q:
                pygame.quit()
                exit()
            elif event.key == pygame.K_l:
                self.game.language = "pt" if self.game.language == "en" else "en"
        return None
        
    def update(self, dt):
        """Update menu (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw menu interface"""
        from game import UIManager, BIG, MED, FONT, WHITE, WIDTH

        title_y = 140
        option_start_y = 300
        option_spacing = 110

        UIManager.draw_text_center(surface, get_text("game_title", self.game.language), BIG, WHITE, title_y)
        UIManager.draw_text_center(surface, get_text("start", self.game.language), MED, WHITE, option_start_y)
        UIManager.draw_text_center(surface, get_text("quit", self.game.language), MED, WHITE, option_start_y + option_spacing)

        lang_name = get_text("portuguese", self.game.language) if self.game.language == "pt" else get_text("english", self.game.language)
        lang_text = f"[L] {get_text('language', self.game.language)}: {lang_name}"
        UIManager.draw_text_center(surface, lang_text, MED, WHITE, option_start_y + option_spacing * 2)
        
        hs_text = f"{get_text('highscore', self.game.language)}: {self.game.highscore.get('name')} - {self.game.highscore.get('score')}"
        hs_surf = FONT.render(hs_text, True, WHITE)
        surface.blit(hs_surf, (WIDTH - hs_surf.get_width() - 20, 40))


class InstructionsScreen(BaseScreen):
    """
    Instructions screen showing how to play
    """
    
    def __init__(self, game):
        """Initialize instructions screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle instructions input"""
        if event.type == pygame.KEYDOWN:
            return "game"  # Start actual game
        return None
        
    def update(self, dt):
        """Update instructions (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw instructions"""
        from game import UIManager, BIG, MED, FONT, WHITE, HEIGHT
        
        UIManager.draw_text_center(surface, get_text("instructions_title", self.game.language), BIG, WHITE, 100)
        UIManager.draw_text_center(surface, get_text("instructions_1", self.game.language), MED, WHITE, 240)
        UIManager.draw_text_center(surface, get_text("instructions_2", self.game.language), MED, WHITE, 320)
        UIManager.draw_text_center(surface, get_text("instructions_continue", self.game.language), FONT, WHITE, HEIGHT - 60)


class GameScreen(BaseScreen):
    """
    Main gameplay screen with enemies, defender, and UI
    """
    
    def __init__(self, game):
        """Initialize game screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle gameplay input"""
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                self.game.pause_background_music()
                return "pause"
            else:
                c = event.unicode
                if c:
                    # Check hits against enemies
                    for enemy in self.game.level_manager.enemies:
                        if enemy.hit(c):
                            break
        return None
        
    def update(self, dt):
        """Update game logic"""
        self.game.level_manager.update()
        self.game.defender.update()
        
        # Update enemies using sprite group
        for enemy in self.game.level_manager.enemies:
            enemy.update(dt)
            enemy.move(dt if enemy.boss else None)
            
            # Check if enemy is dead
            if enemy.is_dead():
                # Create explosion using game_objects.Explosion
                from game_objects import Explosion
                explosion = Explosion(
                    enemy.rect.centerx, enemy.rect.centery,
                    self.game.explosion_frames
                )
                self.game.explosions.add(explosion)
                self.game.explosion_sound.play()
                
                self.game.level_manager.enemies.remove(enemy)
                self.game.score += len(enemy.text) * 5 if not enemy.boss else len(enemy.text) * 8
                self.game.defender.start_shooting()
                continue
                
            # Check collision with tower using sprite collision
            collision_rect = enemy.get_collision_rect() if hasattr(enemy, 'get_collision_rect') else enemy.rect
            if collision_rect.left <= self.game.tower_rect.right:
                # Get damage based on enemy type from config
                dmg = ENEMY_DAMAGE.get(enemy.enemy_type, 1)  # Default to 1 if type not found
                self.game.lives -= dmg
                self.game.ouch_sound.play()
                self.game.level_manager.enemies.remove(enemy)
                
                if self.game.lives <= 0:
                    self.game.stop_background_music()
                    return "lose"
        
        # Check level completion
        lm = self.game.level_manager
        if lm.level_complete:
            lm.level_complete_timer -= 1
            if lm.level_complete_timer <= 0:
                if lm.level < 10:
                    lm.level += 1
                    self.game.level_up_sound.play()
                    lm.start_level_transition()
                else:
                    self.game.stop_background_music()
                    return "congratulations"
                lm.level_complete = False
        
        # Check for boss spawning
        if (not lm.enemies and lm.spawned_count >= self.game.enemy_counts[lm.level] 
            and not lm.transitioning and not lm.awaiting_boss):
            lm.spawn_major_bosses()
        
        # Check for level complete
        if (not lm.enemies and lm.spawned_count >= self.game.enemy_counts[lm.level] 
            and lm.awaiting_boss and not lm.level_complete):
            lm.level_complete = True
            lm.level_complete_timer = 180
                
        return None
        
    def draw(self, surface):
        """Draw game screen"""
        from game import UIManager, BIG, FONT, WHITE, WIDTH, HEIGHT
        
        # Draw defender
        self.game.defender.draw(surface)
        
        # Draw enemies using sprite group
        for enemy in self.game.level_manager.enemies:
            surface.blit(enemy.image, enemy.rect)
            enemy.draw_textbox(surface)
        
        # Draw explosions using sprite group
        self.game.explosions.draw(surface)
        
        # Draw UI elements
        try:
            font_path = os.path.join(os.path.dirname(__file__), "Assets", "Font", "PressStart2P.ttf")
            fonte = pygame.font.Font(font_path, 60)
            img = fonte.render(chr(9829) * self.game.lives, True, (255, 0, 0))
            surface.blit(img, (18, 160))
        except:
            # Fallback if font not found
            heart_text = "♥" * self.game.lives
            heart_surf = FONT.render(heart_text, True, (255, 0, 0))
            surface.blit(heart_surf, (18, 160))
        
        lm = self.game.level_manager
        surface.blit(FONT.render(f"{get_text('level', self.game.language)}: {lm.level}", True, WHITE), (18, 18))
        surface.blit(FONT.render(f"{get_text('score', self.game.language)}: {self.game.score}", True, WHITE), (18, 90))
        
        hs_text = f"{get_text('highscore', self.game.language)}: {self.game.highscore.get('name')} - {self.game.highscore.get('score')}"
        hs_surf = FONT.render(hs_text, True, WHITE)
        surface.blit(hs_surf, (WIDTH - hs_surf.get_width() - 20, 18))

        # Draw level transition
        if lm.transitioning:
            UIManager.draw_text_center(surface, f"{get_text('level', self.game.language).upper()} {lm.level}", BIG, WHITE, 60)
            level_msg = get_text("level_messages", self.game.language).get(lm.level, "")
            UIManager.draw_fade_box(surface, level_msg, lm.transition_alpha)
        
        # Draw level complete message
        if lm.level_complete:
            UIManager.draw_text_center(surface, get_text("level_complete", self.game.language, lm.level), BIG, WHITE, int(HEIGHT / 2 - 50))


class PauseScreen(BaseScreen):
    """
    Pause screen with resume and menu options
    """
    
    def __init__(self, game):
        """Initialize pause screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle pause input"""
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                self.game.resume_background_music()
                return "game"
            elif event.key == pygame.K_m:
                # Reset the game and keep music playing when going to menu
                self.game.reset_game()
                self.game.ensure_background_music()
                return "menu"
        return None
        
    def update(self, dt):
        """Update pause (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw pause interface"""
        from game import UIManager, BIG, MED, WHITE
        
        UIManager.draw_text_center(surface, get_text("paused", self.game.language), BIG, WHITE, 200)
        UIManager.draw_text_center(surface, get_text("resume", self.game.language), MED, WHITE, 330)
        UIManager.draw_text_center(surface, get_text("menu", self.game.language), MED, WHITE, 420)


class LoseScreen(BaseScreen):
    """
    Game over screen with score and name entry for high scores
    """
    
    def __init__(self, game):
        """Initialize lose screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle lose screen input"""
        if self.game.is_new_high:
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    self.game.save_highscore(self.game.entering_name or "Anon", self.game.score)
                    scores = self.game.load_highscores()
                    self.game.highscore = scores[0] if scores else {"name": "None", "score": 0}
                    self.game.is_new_high = False
                    return "highscores"
                elif event.key == pygame.K_BACKSPACE:
                    self.game.entering_name = self.game.entering_name[:-1]
                else:
                    ch = event.unicode
                    if ch.isalnum() and len(self.game.entering_name) < 12:
                        self.game.entering_name += ch
        else:
            if event.type == pygame.KEYDOWN:
                return "highscores"
        return None
        
    def update(self, dt):
        """Update lose screen (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw lose screen"""
        from game import UIManager, BIG, MED, FONT, RED, WHITE, BLUE
        
        UIManager.draw_text_center(surface, get_text("you_died", self.game.language), BIG, RED, 120)
        UIManager.draw_text_center(surface, f"{get_text('score', self.game.language)}: {self.game.score}", MED, WHITE, 240)
        
        if self.game.is_new_high:
            UIManager.draw_text_center(surface, get_text("new_highscore", self.game.language), FONT, WHITE, 320)
            UIManager.draw_text_center(surface, self.game.entering_name + "_", MED, BLUE, 380)
        else:
            UIManager.draw_text_center(surface, get_text("press_key_highscores", self.game.language), MED, WHITE, 360)


class HighscoresScreen(BaseScreen):
    """
    High scores display screen
    """
    
    def __init__(self, game):
        """Initialize highscores screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle highscores input"""
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_m:
                return "menu"
            elif event.key == pygame.K_r:
                return "restart_game"
        return None
        
    def update(self, dt):
        """Update highscores (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw highscores"""
        from game import UIManager, BIG, MED, WHITE, HEIGHT
        
        scores = self.game.get_top_highscores()
        UIManager.draw_highscore_list(surface, scores, get_text("top_10_highscores", self.game.language), BIG, WHITE, 140)
        UIManager.draw_text_center(surface, f"{get_text('restart', self.game.language)}   {get_text('menu', self.game.language)}", MED, WHITE, HEIGHT - 80)


class CongratulationsScreen(BaseScreen):
    """
    Congratulations screen
    """
    
    def __init__(self, game):
        """Initialize congratulations screen"""
        super().__init__(game)
        
    def handle_event(self, event):
        """Handle congratulations input"""
        if self.game.is_new_high:
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    # Save highscore and go to highscores screen
                    self.game.save_highscore(self.game.entering_name or "Winner", self.game.score)
                    scores = self.game.load_highscores()
                    self.game.highscore = scores[0] if scores else {"name": "None", "score": 0}
                    self.game.is_new_high = False
                    return "highscores"
                elif event.key == pygame.K_BACKSPACE:
                    self.game.entering_name = self.game.entering_name[:-1]
                else:
                    ch = event.unicode
                    if ch and ch.isprintable() and len(self.game.entering_name) < 12:
                        self.game.entering_name += ch
        else:
            if event.type == pygame.KEYDOWN:
                return "highscores"
        return None
        
    def update(self, dt):
        """Update congratulations (no dynamic content)"""
        pass
        
    def draw(self, surface):
        """Draw congratulations screen"""
        from game import UIManager, BIG, MED, FONT, GREEN, WHITE, BLUE, HEIGHT

        title_height = BIG.get_height()
        score_height = MED.get_height()
        hint_height = FONT.get_height()

        header_y = 120
        score_y = header_y + title_height + 20
        instructions_y = score_y + score_height + 40
        instructions_y = min(HEIGHT - hint_height - 40, instructions_y)

        UIManager.draw_text_center(surface, get_text("congratulations", self.game.language), BIG, GREEN, header_y)
        UIManager.draw_text_center(surface, f"{get_text('score', self.game.language)}: {self.game.score}", MED, WHITE, score_y)

        # Instructions
        if self.game.is_new_high:
            UIManager.draw_text_center(surface, get_text("new_highscore", self.game.language), FONT, WHITE, instructions_y)
            UIManager.draw_text_center(surface, self.game.entering_name + "_", MED, BLUE, instructions_y + FONT.get_height() + 10)
            UIManager.draw_text_center(surface, get_text("congratulations_hint", self.game.language), FONT, WHITE, instructions_y + FONT.get_height() + MED.get_height() + 20)
        else:
            UIManager.draw_text_center(surface, get_text("press_key_highscores", self.game.language), FONT, WHITE, instructions_y)

